import json
import os
import sys
import socket

from nameko.rpc import rpc

sys.path.append("../")

from utils import dbutil,Utils
from config import Config

class Train:
    name = "model_train"
    @rpc
    def train(self,user_id, model_name, dataset_name, batch_size, learning_rate):
        """
        训练指定模型
        :param user_id: 用户id
        :param model_name: 模型名称
        :param dataset_name: 数据集名称
        :param batch_size: 批处理大小
        :param learning_rate: 学习率
        :return:
        """
        # 获取数据集路径
        dataset_sql = """
                   select dataset_path from app_dataset
                   where user_id='{}' and dataset_name='{}'
               """.format(user_id, dataset_name)
        dataset_path = dbutil.query(dataset_sql)
        if not dataset_path:
            msg = "数据集{}不存在".format(dataset_name)
            status = 404
        else:
            dataset_path = dataset_path[0][0]
            # 获取模型路径、日志路径
            sql = """
                    select status,model_path from app_model
                    where user_id='{}' and model_name='{}'
                    """.format(user_id, model_name)
            data = dbutil.query(sql)
            if not data:
                msg = "模型{}不存在".format(model_name)
                status = 404
            else:
                status = data[0][0]
                if status == 'running':
                    msg = "模型{}训练中".format(model_name)
                    status = 403
                else:
                    batch_size = int(batch_size)
                    learning_rate = float(learning_rate)
                    dataset_path = os.path.join(dataset_path,'')
                    model_path = data[0][1]
                    log_path = os.path.join(data[0][1], model_name + '_train_logs')

                    client = socket.socket()
                    # 连接服务器
                    addr = ('127.0.0.1', 12345)
                    client.connect(addr)
                    data = {
                        'task': 'ext',
                        'mode': 'train',
                        'have_query': 'T',
                        'language': 'Chinese',
                        'batch_size': batch_size,
                        'lr': learning_rate,
                        'bert_data_path': dataset_path,
                        'model_path': model_path,
                        'log_file': log_path,
                        'user_id': user_id,
                        'model_name': model_name
                    }
                    data = json.dumps(data)
                    # 发送数据
                    client.sendall(data.encode('utf-8'))

                    # 接受服务器返回结果
                    result = client.recv(1024).decode()
                    client.close()
                    if result == '200':
                        msg = "开始训练模型{}...".format(model_name)
                        status = 200
                    else:
                        msg = "程序执行错误"
                        status = 500
        response = {
            'msg': msg,
            'status': status
        }

        return json.dumps(response)

    @rpc
    def get_logs(self):

        return 0